import numpy as np
import h5py
from pathlib import Path
from tqdm import tqdm

data_path = Path('/data/IDLab/DigiHealth/training_snapshots/multiple_sr')
chunk_size = 10_000

with h5py.File(data_path / 'train.hdf5', 'a') as f:
    mask = f['y'][:] < 3
    filtered_indices = np.where(mask)[0]
    n_filtered_samples = len(filtered_indices)

    X = f['X']
    X_cwt = f['X_cwt']
    rr = f['rr']
    y = f['y']

    X_filtered = f.require_dataset('X_filtered', shape=(n_filtered_samples, *X.shape[1:]),
                                  maxshape=(None, *X.shape[1:]),
                                  chunks=True, dtype='f')
    X_cwt_filtered = f.require_dataset('X_cwt_filtered', shape=(n_filtered_samples, *X_cwt.shape[1:]),
                                      maxshape=(None, *X_cwt.shape[1:]),
                                      chunks=True, dtype='f')
    rr_filtered = f.require_dataset('rr_filtered', shape=(n_filtered_samples, *rr.shape[1:]),
                                    maxshape=(None, *rr.shape[1:]),
                                    chunks=True, dtype='f')
    y_filtered = f.require_dataset('y_filtered', shape=(n_filtered_samples, *y.shape[1:]),
                                  maxshape=(None,),
                                  chunks=True, dtype='i')
    
    for start in tqdm(range(0, n_filtered_samples, chunk_size)):
        end = min(start + chunk_size, n_filtered_samples)
        X_filtered[start:end] = X[filtered_indices[start:end]]
        X_cwt_filtered[start:end] = X_cwt[filtered_indices[start:end]]
        rr_filtered[start:end] = rr[filtered_indices[start:end]]
        y_filtered[start:end] = y[filtered_indices[start:end]]